{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "expected-tunisia",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "import numpy as np\n",
    "%matplotlib widget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "extensive-candidate",
   "metadata": {},
   "outputs": [],
   "source": [
    "points_3d = np.array([[ 0.,          0.,          0.        ],\n",
    " [-0.03138264, -0.09922726,  0.00946505],\n",
    " [-0.07037588,  0.07883613,  0.0089331 ],\n",
    " [ 0.10787677,  0.02588217,  0.00836965]])\n",
    "keypoints_2d = np.array([[603.39123535, 301.11212158],\n",
    " [640.75494385, 396.78601074],\n",
    " [556.61352539, 291.11450195],\n",
    " [697.81604004, 276.12750244]])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "alleged-clerk",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "77ed4fe527794923953ac0c2b3680b47",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "(-0.15, 0.15)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(projection='3d')\n",
    "ax.scatter(points_3d[:, 0], points_3d[:, 1], points_3d[:, 2], c=np.linspace(0, 1, 4), cmap='summer')\n",
    "ax.set_xlim(-0.15, 0.15)\n",
    "ax.set_ylim(-0.15, 0.15)\n",
    "ax.set_zlim(-0.15, 0.15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "settled-township",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "326ba677018f4daf96840f6b2f74e98b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "(0.0, 720.0)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "figure = plt.figure()\n",
    "plt.scatter(keypoints_2d[:, 0], keypoints_2d[:, 1], c=np.arange(4), cmap='summer')\n",
    "plt.xlim(0, 1280)\n",
    "plt.ylim(0, 720)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "religious-wells",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "K = np.array([[697.87732212 , 0. , 648.08562626],\n",
    " [  0., 697.28594061, 371.49958099],\n",
    " [  0., 0.,1.        ]]) / 8.0\n",
    "D = np.array([-1.74610270e-01,  2.75427408e-02,  6.24873971e-05,  9.10956548e-05])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "sealed-cambridge",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(True,\n",
       " array([[ 0.75168222, -0.27787168,  0.5981314 , -0.0314786 ],\n",
       "        [-0.45333637, -0.87638764,  0.16257562, -0.06893673],\n",
       "        [ 0.47901981, -0.39335992, -0.78473435,  0.89141784],\n",
       "        [ 0.        ,  0.        ,  0.        ,  1.        ]]))"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "success, rotvec, tvec = cv2.solvePnP(points_3d, keypoints_2d, K, np.zeros(4), flags=cv2.SOLVEPNP_EPNP)\n",
    "T = np.eye(4)\n",
    "R, _ = cv2.Rodrigues(rotvec)\n",
    "T[:3, :3] = R\n",
    "T[:3, 3] = tvec[:, 0]\n",
    "success, T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "regional-gossip",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[-0.0314786 , -0.06893673,  0.89141784,  1.        ],\n",
       "       [-0.02183458,  0.03379049,  0.90798941,  1.        ],\n",
       "       [-0.10094206, -0.10467149,  0.81968532,  1.        ],\n",
       "       [ 0.04742468, -0.1391633 ,  0.92634399,  1.        ]])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predicted_points = (T @ np.concatenate([points_3d, np.ones((4, 1))], axis=1)[:,:, None])[:, :, 0]\n",
    "predicted_points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "valid-folks",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dc3df26192f6474da4c095a2f249e203",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[623.44151511 317.57581015   1.        ]\n",
      " [631.30364683 397.44882302   1.        ]\n",
      " [562.14389836 282.4581459    1.        ]\n",
      " [683.81383124 266.74733235   1.        ]]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(0.0, 720.0)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fig = plt.figure()\n",
    "reprojected = (K @ predicted_points[:, :3, None])[:, :, 0]\n",
    "reprojected = reprojected / reprojected[:, 2:3]\n",
    "print(reprojected)\n",
    "\n",
    "plt.scatter(keypoints_2d[:, 0], keypoints_2d[:, 1], c=np.linspace(0, 1, keypoints_2d.shape[0]), cmap='spring')\n",
    "plt.scatter(reprojected[:, 0], reprojected[:, 1], c=np.linspace(0, 1, reprojected.shape[0]), cmap='spring')\n",
    "plt.xlim(0, 1280)\n",
    "plt.ylim(0, 720)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "formal-marketplace",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[525.72722885, 262.88184879,   0.83918927],\n",
       "       [564.64292418, 344.10913791,   0.88081192],\n",
       "       [450.11823921, 233.91501349,   0.82930932],\n",
       "       [543.28804458, 205.62918175,   0.78049184]])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(K @ predicted_points[:, :3, None])[:, :, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "permanent-assessment",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[-0.02599206, -0.07009551,  0.83918927,  1.        ],\n",
       "       [-0.0088821 ,  0.02421945,  0.88081192,  1.        ],\n",
       "       [-0.1251584 , -0.10637394,  0.82930932,  1.        ],\n",
       "       [ 0.05368064, -0.1209306 ,  0.78049184,  1.        ]])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predicted_points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "mature-corpus",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "vision",
   "language": "python",
   "name": "vision"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
